package com.spark.common.utils;

import com.spark.common.core.entity.model.DataSourceColumnsModel;
import com.spark.common.core.entity.model.DataSourceTablesModel;
import com.spark.common.exception.SparkException;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.CannotGetJdbcConnectionException;

import java.sql.*;
import java.util.ArrayList;
import java.util.List;

/**
 * 数据源工具类
 *
 * @author LYCHEE
 * @date 2025/2/18 10:28
 */
@Data
@Slf4j
public class DataSourceUtils {

    private String dbDriver;
    private String dbUrl;
    private String dbUsername;
    private String dbPassword;

    public DataSourceUtils(String dbDriver, String dbUrl, String dbUsername, String dbPassword) {
        this.dbDriver = dbDriver;
        this.dbUrl = dbUrl;
        this.dbUsername = dbUsername;
        this.dbPassword = dbPassword;
    }

    /**
     * 测试连接
     */
    public void testConnect() {
        Connection connection = null;
        try {
            connection = getConnection(dbDriver, dbUrl, dbUsername, dbPassword);
        } catch (Exception e) {
            throw new SparkException("数据库连接失败：" + e.getMessage());
        } finally {
            try {
                if (connection != null && !connection.isClosed()) {
                    connection.close();
                }
            } catch (SQLException e) {
                log.error(e.getMessage());
            }
        }
    }

    /**
     * 查询数据库中的所有表信息
     *
     * @return 表信息
     */
    public List<DataSourceTablesModel> queryAllTablesInfo() {
        List<DataSourceTablesModel> tablesList = new ArrayList<>();
        Connection connection = null;
        try {
            connection = getConnection(dbDriver, dbUrl, dbUsername, dbPassword);
            DatabaseMetaData metaData = connection.getMetaData();
            ResultSet resultSet = metaData.getTables(extractDatabaseNameFromJdbcUrl(dbUrl), null, "%", new String[]{"TABLE"});
            while (resultSet.next()) {
                DataSourceTablesModel tableInfo = new DataSourceTablesModel();
                // 目录名称
                tableInfo.setTableCat(resultSet.getString("TABLE_CAT"));
                // 模式名称
                tableInfo.setTableSchem(resultSet.getString("TABLE_SCHEM"));
                // 表名称
                tableInfo.setTableName(resultSet.getString("TABLE_NAME"));
                // 表类型
                tableInfo.setTableType(resultSet.getString("TABLE_TYPE"));
                // 表注释
                String remark = getTableComments(connection, tableInfo.getTableName());
                tableInfo.setRemark(remark);
                tablesList.add(tableInfo);
            }
        } catch (Exception e) {
            throw new SparkException("数据库连接失败：" + e.getMessage());
        } finally {
            try {
                if (connection != null && !connection.isClosed()) {
                    connection.close();
                }
            } catch (SQLException e) {
                log.error(e.getMessage());
            }
        }
        return tablesList;
    }

    /**
     * 查询数据库中的指定表的字段信息
     *
     * @param tableName 表名
     * @return 字段信息
     */
    public List<DataSourceColumnsModel> queryTableColumnInfo(String tableName) {
        List<DataSourceColumnsModel> columnInfoList = new ArrayList<>();
        Connection connection = null;
        PreparedStatement columnsStmt = null;
        PreparedStatement foreignKeysStmt = null;
        PreparedStatement uniqueKeysStmt = null;
        try {
            connection = getConnection(dbDriver, dbUrl, dbUsername, dbPassword);
            String columnsQuery = "SELECT " +
                    "COLUMN_NAME AS '字段名称', " +
                    "COLUMN_COMMENT AS '字段注释', " +
                    "DATA_TYPE AS '字段类型', " +
                    "CHARACTER_MAXIMUM_LENGTH AS '字段长度', " +
                    "NUMERIC_SCALE AS '小数点位数', " +
                    "COLUMN_DEFAULT AS '默认值', " +
                    "CASE WHEN COLUMN_KEY = 'PRI' THEN '是' ELSE '否' END AS '是否主键', " +
                    "CASE WHEN IS_NULLABLE = 'YES' THEN '是' ELSE '否' END AS '允许NULL' " +
                    "FROM INFORMATION_SCHEMA.COLUMNS " +
                    "WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?";
            columnsStmt = connection.prepareStatement(columnsQuery);
            columnsStmt.setString(1, connection.getCatalog());
            columnsStmt.setString(2, tableName);
            ResultSet columnsRs = columnsStmt.executeQuery();
            while (columnsRs.next()) {
                DataSourceColumnsModel columnsModel = new DataSourceColumnsModel(
                        columnsRs.getString("字段名称"),
                        columnsRs.getString("字段注释"),
                        columnsRs.getString("字段类型"),
                        columnsRs.getInt("字段长度"),
                        columnsRs.getInt("小数点位数"),
                        columnsRs.getString("默认值"),
                        columnsRs.getString("是否主键"),
                        columnsRs.getString("允许NULL"),
                        null,
                        null,
                        null
                );
                columnInfoList.add(columnsModel);
            }
            // 查询外键信息
            String foreignKeysQuery = "SELECT " +
                    "REFERENCED_TABLE_NAME AS '外键表', " +
                    "COLUMN_NAME AS '字段名称', " +
                    "REFERENCED_COLUMN_NAME AS '外键字段' " +
                    "FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE " +
                    "WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? AND REFERENCED_TABLE_NAME IS NOT NULL";
            foreignKeysStmt = connection.prepareStatement(foreignKeysQuery);
            foreignKeysStmt.setString(1, connection.getCatalog());
            foreignKeysStmt.setString(2, tableName);
            ResultSet foreignKeysRs = foreignKeysStmt.executeQuery();

            while (foreignKeysRs.next()) {
                String columnName = foreignKeysRs.getString("字段名称");
                for (DataSourceColumnsModel columnsModel : columnInfoList) {
                    if (columnsModel.getColumnName().equals(columnName)) {
                        columnsModel.setForeignKeyTable(foreignKeysRs.getString("外键表"));
                        columnsModel.setForeignKeyField(foreignKeysRs.getString("外键字段"));
                        break;
                    }
                }
            }
            // 查询唯一约束信息
            String uniqueKeysQuery = "SELECT " +
                    "KCU.COLUMN_NAME AS '字段名称', " +
                    "TC.CONSTRAINT_TYPE AS '约束类型' " +
                    "FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS TC " +
                    "JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE KCU ON TC.CONSTRAINT_NAME = KCU.CONSTRAINT_NAME " +
                    "WHERE TC.TABLE_SCHEMA = ? AND TC.TABLE_NAME = ? AND TC.CONSTRAINT_TYPE = 'UNIQUE'";
            uniqueKeysStmt = connection.prepareStatement(uniqueKeysQuery);
            uniqueKeysStmt.setString(1, connection.getCatalog());
            uniqueKeysStmt.setString(2, tableName);
            ResultSet uniqueKeysRs = uniqueKeysStmt.executeQuery();
            while (uniqueKeysRs.next()) {
                String columnName = uniqueKeysRs.getString("字段名称");
                for (DataSourceColumnsModel columnsModel : columnInfoList) {
                    if (columnsModel.getColumnName().equals(columnName)) {
                        columnsModel.setIsUnique("是");
                        break;
                    }
                }
            }
        } catch (Exception e) {
            throw new SparkException("数据库连接失败：" + e.getMessage());
        } finally {
            try {
                if (connection != null && !connection.isClosed()) {
                    connection.close();
                }
                if (columnsStmt != null) {
                    columnsStmt.close();
                }
                if (foreignKeysStmt != null) {
                    foreignKeysStmt.close();
                }
                if (uniqueKeysStmt != null) {
                    foreignKeysStmt.close();
                }
            } catch (SQLException e) {
                log.error(e.getMessage());
            }
        }
        return columnInfoList;
    }

    /**
     * 判断数据库中是否存在某个表
     *
     * @param tableName 表名
     * @return 是否存在
     */
    public boolean existsTable(String tableName) {
        Connection connection = null;
        Statement statement = null;
        try {
            connection = getConnection(dbDriver, dbUrl, dbUsername, dbPassword);
            statement = connection.createStatement();
            // 获取数据库表名
            ResultSet tables = connection.getMetaData().getTables(null, null, tableName, null);
            if (tables.next()) {
                return true;
            }
        } catch (Exception e) {
            throw new SparkException(e.getMessage());
        } finally {
            try {
                if (connection != null && !connection.isClosed()) {
                    connection.close();
                }
                if (statement != null) {
                    statement.close();
                }
            } catch (SQLException e) {
                log.error(e.getMessage());
            }
        }
        return false;
    }

    /**
     * 删除数据库中的表
     *
     * @param tableName 表名
     * @return 删除状态
     */
    public boolean deleteTable(String tableName) {
        Connection connection = null;
        Statement statement = null;
        try {
            connection = getConnection(dbDriver, dbUrl, dbUsername, dbPassword);
            statement = connection.createStatement();
            String dropTableSql = "DROP TABLE IF EXISTS `" + tableName + "`";
            statement.executeUpdate(dropTableSql);
            return true;
        } catch (Exception e) {
            log.error("删除表失败：{}", e.getMessage());
            return false;
        } finally {
            try {
                if (connection != null && !connection.isClosed()) {
                    connection.close();
                }
                if (statement != null) {
                    statement.close();
                }
            } catch (SQLException e) {
                log.error(e.getMessage());
            }
        }
    }

    /**
     * 执行SQL
     *
     * @param sql        SQL语句
     */
    public boolean executeSql(String sql) {
        Connection connection = null;
        Statement statement = null;
        try {
            // 获取连接
            connection = getConnection(dbDriver, dbUrl, dbUsername, dbPassword);
            statement = connection.createStatement();
            statement.execute(sql);
            return true;
        } catch (Exception e) {
            log.error("执行SQL语句失败：{}", e.getMessage());
            return false;
        } finally {
            try {
                if (connection != null && !connection.isClosed()) {
                    connection.close();
                }
                if (statement != null) {
                    statement.close();
                }
            } catch (SQLException e) {
                log.error(e.getMessage());
            }
        }
    }

    /**
     * 获取数据库连接
     *
     * @param dbDriver   驱动类
     * @param dbUrl      连接Url
     * @param dbUsername 用户名
     * @param dbPassword 密码
     * @return 返回连接
     */
    private Connection getConnection(String dbDriver, String dbUrl, String dbUsername, String dbPassword) {
        Connection connection = null;
        try {
            Class.forName(dbDriver);
            connection = DriverManager.getConnection(dbUrl, dbUsername, dbPassword);
            if (connection == null) {
                throw new SparkException("数据库连接失败，获取的连接为空或已关闭");
            }
        } catch (SQLException e) {
            if (e.getMessage().toLowerCase().contains("access denied") || e.getMessage().toLowerCase().contains("bad password")) {
                throw new SparkException("数据库连接失败：用户名或密码错误");
            } else {
                throw new SparkException("数据库连接失败：" + e.getMessage());
            }
        } catch (CannotGetJdbcConnectionException e) {
            throw new SparkException("数据库连接失败,无法获取JDBC连接：" + e.getMessage());
        } catch (Exception e) {
            throw new SparkException("数据库连接失败：" + e.getMessage());
        }
        return connection;
    }

    /**
     * 从JDBC URL中提取数据库名称
     *
     * @param jdbcUrl JDBC连接URL
     * @return 数据库名称
     */
    private String extractDatabaseNameFromJdbcUrl(String jdbcUrl) {
        String regex = "^jdbc:mysql://[^/]+/([^?]+)";
        java.util.regex.Pattern pattern = java.util.regex.Pattern.compile(regex);
        java.util.regex.Matcher matcher = pattern.matcher(jdbcUrl);
        if (matcher.find()) {
            return matcher.group(1);
        } else {
            throw new IllegalArgumentException("No database name found in the JDBC URL");
        }
    }

    /**
     * 获取指定表的注释
     *
     * @param connection 数据库连接
     * @param tableName  表名称
     * @return 表的注释
     * @throws SQLException SQL异常
     */
    private String getTableComments(Connection connection, String tableName) throws SQLException {
        String query = "SELECT TABLE_COMMENT FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?";
        try (PreparedStatement preparedStatement = connection.prepareStatement(query)) {
            // 使用当前连接的数据库名称
            preparedStatement.setString(1, connection.getCatalog());
            preparedStatement.setString(2, tableName);
            try (ResultSet rs = preparedStatement.executeQuery()) {
                if (rs.next()) {
                    return rs.getString("TABLE_COMMENT");
                }
            }
        }
        // 如果没有找到注释，则返回null
        return null;
    }
}
